"""文本处理工具 - 基于 Token 的文本分块"""
import tiktoken
from typing import List, Dict


class TextProcessor:
    """文本处理器"""
    
    def __init__(self, chunk_size: int = 500):
        """
        初始化文本处理器
        
        Args:
            chunk_size: 每个文本块的 token 数量
        """
        self.chunk_size = chunk_size
        # 使用与 text-embedding-3-small 匹配的编码器
        self.encoding = tiktoken.get_encoding("cl100k_base")
    
    def chunk_text(self, text: str) -> List[Dict[str, any]]:
        """
        按 token 数量分块文本
        
        Args:
            text: 待分块的文本
            
        Returns:
            包含分块信息的列表，每个元素包含 content 和 token_count
        """
        # 编码文本为 token
        tokens = self.encoding.encode(text)
        
        # 按固定大小分割
        chunks = []
        for i in range(0, len(tokens), self.chunk_size):
            chunk_tokens = tokens[i:i + self.chunk_size]
            # 解码回文本
            chunk_text = self.encoding.decode(chunk_tokens)
            chunks.append({
                "content": chunk_text,
                "token_count": len(chunk_tokens),
                "index": len(chunks)
            })
        
        return chunks
    
    def count_tokens(self, text: str) -> int:
        """
        计算文本的 token 数量
        
        Args:
            text: 待计算的文本
            
        Returns:
            token 数量
        """
        return len(self.encoding.encode(text))

